from torch_geometric.loader import LinkNeighborLoader,ZipLoader
import torch
import copy
import torch_geometric.transforms as T
from data_package_003 import my_dataset
from torch_geometric.loader import DataLoader
from torch_geometric.data import Data

def descrition_data(_data_dict, _data):
    _descrition_data = {}
    all_list = []
    vdata = _data.edge_label_index.t().tolist()
    # print(vdata)
    for i2j in vdata:
        all_list.append((i2j, _data_dict.all_dict[i2j[0]], _data_dict.all_dict[i2j[1]]))
    return all_list


transform = T.RandomLinkSplit(
    num_val=0.2,
    num_test=0.1,
    disjoint_train_ratio=0.1,
    neg_sampling_ratio=2.0,
    add_negative_train_samples=False,
    is_undirected=True,
#    split_labels=True,
    edge_types=[
        # ("alarm", "on", "host"),
        # ("alarm", "to", "btree"),
        ("host", "belongsto", "btree")
    ],
)

# <class 'torch_geometric.data.hetero_data.HeteroData'>
# <class 'torch_geometric.data.hetero_data.HeteroData'>
train_data_list = []

#
# print(next(iter(my_dataset)))
datas = []

edge_index_list = []
node_id_list = []
x_list = []
node_type_list = []
edge_type_list = []

"""
Data(edge_index=[2, 1136], node_id=[389], x=[389, 768], node_type=[389], edge_type=[1136])
Data(edge_index=[2, 948], node_id=[389], x=[389, 768], node_type=[389], edge_type=[948])
Data(edge_index=[2, 1124], node_id=[389], x=[389, 768], node_type=[389], edge_type=[1124])
Data(edge_index=[2, 1170], node_id=[367], x=[367, 768], node_type=[367], edge_type=[1170])
"""

first = True
for batch in my_dataset:
    print(batch,"2222")
    data = batch.to_homogeneous()
    edge_index_list.append(data.edge_index)
    edge_type_list.append(data.edge_type)
    if first == True:
        first = False
        node_id_list.append(data.node_id[data.node_type==0])
        node_id_list.append(data.node_id[data.node_type==1])
        
        x_list.append(data.x[data.node_type==0])
        x_list.append(data.x[data.node_type==1])

        node_type_list.append(data.node_type[data.node_type==0])
        node_type_list.append(data.node_type[data.node_type==1])

    node_id_list.append(data.node_id[data.node_type==2])
    x_list.append(data.x[data.node_type==2])
    node_type_list.append(data.node_type[data.node_type==2])
    #print(data.node_id[data.node_type==0])
    #print(data.node_id[data.node_type==1])
    #print(data.node_id[data.node_type==2])

#print(torch.cat(edge_index_list,dim=1).shape)
#print(torch.cat(node_id_list,dim=0).shape)
#print(torch.cat(x_list,dim=0).shape)
#print(torch.cat(node_type_list,dim=0).shape)
#print(torch.cat(edge_type_list,dim=0).shape)

origin_data = my_dataset.get(0)

data_sum = Data(edge_index = torch.cat(edge_index_list,dim=1) ,node_id= torch.cat(node_id_list,dim=0),x = torch.cat(x_list,dim=0),node_type=torch.cat(node_type_list,dim=0),edge_type = torch.cat(edge_type_list,dim=0)) 
# print(data_sum)
train_data, val_data, test_data = transform(data_sum)
# print(train_data.edge_index)

train_loader = LinkNeighborLoader(
        data=train_data,
        num_neighbors=[10, 10],
        neg_sampling_ratio=5.0,
        edge_label_index=train_data.edge_label_index.to(torch.long),
        edge_label=train_data.edge_label.to(torch.long),
        batch_size=16,
        shuffle=True,
        drop_last=True,
    )
#    print("train_loadertrain_loadertrain_loader",train_loader)
#    train_data_list.append(train_loader)

#
# # print("train_data",train_data.node_id)
# # print("val_data",val_data)
# # print("test_data",test_data)
#
# s =  next(iter(train_loader))
# print("1111111",s.x)